''' nn.Conv2d '''
from torch import nn
import mymnist
import mytb

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(1, 1, kernel_size=3, padding=1) # 卷积层

    def forward(self, x):
        x = self.conv(x)
        return x

with mytb.tb() as writer:
    for step, (images, labels) in enumerate(mymnist.test_loader):
        writer.add_images('input', images, step)
        output = Net()(images)
        writer.add_images('output', output, step)
